#include <torch/extension.h>
#include <vector>

__global__ void render_kernel(
    const float* __restrict__ point2D,
    const float* __restrict__ color,
    const float* __restrict__ opacity,
    const float* __restrict__ depths,
    float* render_color,
    float* render_depth,
    float* render_alpha,
    const int N,
    const int H,
    const int W) {

    int h = blockIdx.y * blockDim.y + threadIdx.y;
    int w = blockIdx.x * blockDim.x + threadIdx.x;

    if (h >= H || w >= W) return;

    float back_ground[3] = {1.0, 1.0, 1.0};
    float pixel_color[3] = {1.0, 1.0, 1.0};
    float pixel_depth = 0.0;
    float pixel_alpha = 0.0;

    for (int i = 0; i < N; ++i) {
        float point_h = point2D[i * 2 + 0];
        float point_w = point2D[i * 2 + 1];

        if (point_h >= h && point_h < h + 1 && point_w >= w && point_w < w + 1) {
            float depth = depths[i];
            float alpha = opacity[i] * (1.0 - pixel_alpha);

            for (int c = 0; c < 3; ++c) {
                pixel_color[c] = pixel_color[c] * (1.0 - alpha) + color[i * 3 + c] * alpha;
            }
            pixel_depth = pixel_depth * (1.0 - alpha) + depth * alpha;
            pixel_alpha += alpha;
        }
    }

    for (int c = 0; c < 3; ++c) {
        pixel_color[c] = (1.0 - pixel_alpha) * back_ground[c] + pixel_alpha * pixel_color[c];
    }

    int pixel_idx = h * W + w;
    for (int c = 0; c < 3; ++c) {
        render_color[pixel_idx * 3 + c] = pixel_color[c];
    }
    render_depth[pixel_idx] = pixel_depth;
    render_alpha[pixel_idx] = pixel_alpha;
}

std::vector<torch::Tensor> render(
    torch::Tensor point2D,
    torch::Tensor color,
    torch::Tensor opacity,
    torch::Tensor depths,
    int H,
    int W) {

    const int threads = 16;
    const dim3 blocks((W + threads - 1) / threads, (H + threads - 1) / threads);
    const dim3 threads_per_block(threads, threads);

    auto render_color = torch::ones({H * W * 3}, point2D.options());
    auto render_depth = torch::zeros({H * W}, point2D.options());
    auto render_alpha = torch::zeros({H * W}, point2D.options());

    render_kernel<<<blocks, threads_per_block>>>(
        point2D.data_ptr<float>(),
        color.data_ptr<float>(),
        opacity.data_ptr<float>(),
        depths.data_ptr<float>(),
        render_color.data_ptr<float>(),
        render_depth.data_ptr<float>(),
        render_alpha.data_ptr<float>(),
        point2D.size(0), H, W);

    return {render_color.view({H, W, 3}), render_depth.view({H, W, 1}), render_alpha.view({H, W, 1})};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("render", &render, "Render kernel");
}
